#include "avs4x264mod.h"

extern avs_hnd_t avs_h;
extern buffer_t OBuffer;

char* GenerateCommand(int argc, char ** argv, video_info_t * VideoInfo, pipe_info_t * PipeInfo)
{
	int i;
	char *cmd, *buf;
	int b_tc = 0;
	int len = (unsigned int)strrchr(argv[0], '\\');
	char * x264_binary;
	x264_binary = "x264_64.exe";
	if (len)
	{
		len -=(unsigned int)argv[0];
		buf = malloc(len+2<20 ? 20: len+2);
		strncpy(buf, argv[0], len);
		buf[len] = '\\';
		buf[len+1] = 0;
	}
	else
	{
		buf = malloc(8192);
		*buf=0;
	}
	cmd = malloc(8192);
	for (i=1;i<argc;i++)
	{
		if( !strcmp(argv[i], "--x264-binary") || !strcmp(argv[i], "-L") )
		{
			x264_binary = argv[i+1];
			memmove(argv + i, argv + i + 2, (argc - i - 2) * sizeof(char *));
			argc -= 2;
			i--;
		}
		else if(!strcmp(argv[i], "--affinity"))
		{
			PipeInfo->Affinity = atoi(argv[i+1]);
			memmove(argv + i, argv + i + 2, (argc - i - 2) * sizeof(char *));
			argc -= 2;
			i--;
		}
		else if(!strcmp(argv[i], "--x264-affinity"))
		{
			PipeInfo->x264Affinity = atoi(argv[i+1]);
			memmove(argv + i, argv + i + 2, (argc - i - 2) * sizeof(char *));
			argc -= 2;
			i--;
		}
		else if(!strcmp(argv[i], "--pipe-buffer"))
		{
			VideoInfo->b_mtpipe = 1;
			PipeInfo->BufferSize = atoi(argv[i+1]);
			memmove(argv + i, argv + i + 2, (argc - i - 2) * sizeof(char *));
			argc -= 2;
			i--;
		}
		else if(!strcmp(argv[i], "--save-log"))
		{
			PipeInfo->SaveLog = argv[i+1];
			memmove(argv + i, argv + i + 2, (argc - i - 2) * sizeof(char *));
			argc -= 2;
			i--;
		}
		else if(!strcmp(argv[i], "--frames"))
		{
			PipeInfo->Frames = atoi(argv[i+1]);
			memmove(argv + i, argv + i + 2, (argc - i - 2) * sizeof(char *));
			argc -= 2;
			i--;
		}
		else if(!strcmp(argv[i], "--seek"))
		{
			PipeInfo->Seek = atoi(argv[i+1]);
			memmove(argv + i, argv + i + 2, (argc - i - 2) * sizeof(char *));
			argc -= 2;
			i--;
		}
	}
	for (i=1;i<argc;i++)
	{
		if( !strcmp(argv[i], "--pipe-mt") )
		{
			VideoInfo->b_mtpipe = 1;
			memmove(argv + i, argv + i + 1, (argc - i - 1) * sizeof(char *));
			argc--;
			break;
		}
	}
	for (i=1;i<argc;i++)
	{
		if( !strcmp(argv[i], "--tcfile-in") )
			b_tc = 1;
		else if( !strcmp(argv[i], "--input-depth") )
		{
			if( strcmp(argv[++i], "8") )
				VideoInfo->i_width /= 2;
		}
	}
	if(PipeInfo->Frames == 0 || 	PipeInfo->Frames > VideoInfo->i_frame_total - PipeInfo->Seek)
		PipeInfo->Frames = VideoInfo->i_frame_total - PipeInfo->Seek;
	if( b_tc )
		sprintf(cmd, "%s%s\" --frames %d ", buf, x264_binary, PipeInfo->Frames);
	else
		sprintf(cmd, "%s%s\" --frames %d --fps %d/%d ", buf, x264_binary, PipeInfo->Frames, VideoInfo->i_fps_num, VideoInfo->i_fps_den);

	/* skip invalid path name when both avs4x264mod and x264 binary is given by full path */
	int p_cmd = (int)cmd;
	char *cmd_tmp;
	cmd_tmp = malloc(8192);
	int cmd_len = strlen(cmd);
	cmd = strrchr( cmd, '\"' );											  /* find the end of x264 binary path */
	while( strlen(cmd) < cmd_len && *(cmd) != ':' )						  /* find if drive number is given */
		cmd--;
	while( strlen(cmd--) < cmd_len && *(cmd) != '\\' && *(cmd) != '/' );	 /* if find drive number, skip invalid path before it */
	*cmd = '"';															   /* insert '"' before processed path */
	strcpy(cmd_tmp, cmd);
	cmd = (char *)p_cmd;
	strcpy(cmd, cmd_tmp);
	free(cmd_tmp);
	
	for (i=1;i<argc;i++)
	{
		if (VideoInfo->infile!=argv[i])
		{
			if (strrchr(argv[i], ' '))
			{
				strcat(cmd, "\"");
				strcat(cmd, argv[i]);
				strcat(cmd, "\" ");
			}
			else
			{
				strcat(cmd, argv[i]);
				strcat(cmd, " ");
			}
		}
	}
	sprintf(buf, "- --input-res %dx%d", VideoInfo->i_width, VideoInfo->i_height);
	strcat(cmd, buf);
	free(buf);
	return cmd;
}

void __cdecl BufferedPipe(void * p)
{
	int result;
	DWORD temp;
	pipe_info_t * PipeInfo = (pipe_info_t *)p;
	SetThreadPriority(GetCurrentThread(), THREAD_PRIORITY_TIME_CRITICAL);
	while(1)
	{
		while(OBuffer.size == 0)
		{
			//printf("~%d~", OBuffer.size);
			usleep(100000);
		}
		if(OBuffer.size < 0)
			return;
		node_t* clink = OBuffer.dynarray+OBuffer.head;
		if(clink->length > 0)
		{
			result = WriteFile(PipeInfo->h_pipeWrite, clink->data, clink->length, &temp, NULL);
			if(!result)
			{
				//printf("~%d\n", clink->length);
				//putchar('&');
				OBuffer.exitcode = -1;
				return;
			}
		}
		//printf("-%d*%d ", clink->length, OBuffer.head);
		free(clink->data);
		clink->data = 0;
		clink->length = 0;
		OBuffer.head++;
		if(OBuffer.head >= OBuffer.capacity)
			OBuffer.head -= OBuffer.capacity;
		OBuffer.size--;
	}
}

#define DATASIZE 1048576
void __cdecl BufferedReadErrorPipe(void * p)
{
	int result;
	DWORD temp;
	pipe_info_t * PipeInfo = (pipe_info_t *)p;
	SetThreadPriority(GetCurrentThread(), THREAD_PRIORITY_TIME_CRITICAL);
	FILE * fp = NULL;
	if(PipeInfo->SaveLog)
		fp = fopen(PipeInfo->SaveLog, "a+");
	char data[DATASIZE];
	while(1)
	{
		result = ReadFile(PipeInfo->h_errorRead, data, DATASIZE, &temp, NULL);
		if(!result)
		{
			//if(GetLastError() == ERROR_HANDLE_EOF)
				//return;
			break;
		}
		if(data[temp - 1] != '\r' && fp)
			fwrite(data, 1, sizeof(data), fp);
		data[temp - 1] = 0;
		fprintf(stderr, "%s\tBuffer: %d/%d (%d MB)\r", data, OBuffer.size, OBuffer.capacity, (OBuffer.framesize * OBuffer.size) >> 20);
		fflush(stderr);
	}
	if(fp)
		fclose(fp);
}

int CreateX264Process(int argc, char** argv, video_info_t * VideoInfo, pipe_info_t * PipeInfo)
{
	char *cmd;
	//createprocess related
	//execute the commandline
	PipeInfo->h_process = GetCurrentProcess();
	PipeInfo->h_stdOut = GetStdHandle(STD_OUTPUT_HANDLE);
	//PipeInfo->h_stdErr = GetStdHandle(STD_ERROR_HANDLE);

	if (PipeInfo->h_stdOut==INVALID_HANDLE_VALUE) // || PipeInfo->h_stdErr==INVALID_HANDLE_VALUE
	{
		fprintf( stderr, "Error: Couldn\'t get standard handles!");
		return ERR_AVS_FAIL;
	}
	
	PipeInfo->saAttr.nLength = sizeof(SECURITY_ATTRIBUTES);
	PipeInfo->saAttr.bInheritHandle = TRUE;
	PipeInfo->saAttr.lpSecurityDescriptor = NULL;

	if (!CreatePipe(&PipeInfo->h_pipeRead, &PipeInfo->h_pipeWrite, &PipeInfo->saAttr, PIPE_BUFFER_SIZE))
	{
		fprintf( stderr, "Error: Pipe creation failed!");
		return ERR_AVS_FAIL;
	}

	if (!CreatePipe(&PipeInfo->h_errorRead, &PipeInfo->h_errorWrite, &PipeInfo->saAttr, PIPE_BUFFER_SIZE))
	{
		fprintf( stderr, "Error: Pipe2 creation failed!");
		return ERR_AVS_FAIL;
	}

	if ( !SetHandleInformation(PipeInfo->h_pipeWrite, HANDLE_FLAG_INHERIT, 0) )
	{
		fprintf( stderr, "Error: SetHandleInformation");
		return ERR_PIPE_FAIL;
	}

	if ( !SetHandleInformation(PipeInfo->h_errorRead, HANDLE_FLAG_INHERIT, 0) )
	{
		fprintf( stderr, "Error: SetHandleInformation");
		return ERR_PIPE_FAIL;
	}
	
	ZeroMemory( &PipeInfo->pi_info, sizeof(PROCESS_INFORMATION) );
	ZeroMemory( &PipeInfo->si_info, sizeof(STARTUPINFO) );
	PipeInfo->si_info.cb = sizeof(STARTUPINFO);
	PipeInfo->si_info.dwFlags = STARTF_USESTDHANDLES;
	PipeInfo->si_info.hStdInput = PipeInfo->h_pipeRead;
	PipeInfo->si_info.hStdOutput = PipeInfo->h_stdOut;
	PipeInfo->si_info.hStdError = PipeInfo->h_errorWrite;
	
	cmd = GenerateCommand(argc, argv, VideoInfo, PipeInfo);
	fprintf( stderr, "avs4x264 [info]: %s\n", cmd);
	fflush(stderr);

	if (!CreateProcess(NULL, cmd, NULL, NULL, TRUE, 0, NULL, NULL, &PipeInfo->si_info, &PipeInfo->pi_info))
	{
		fprintf( stderr, "Error: Failed to create process <%d>!", (int)GetLastError());
		free(cmd);
		return ERR_PIPE_FAIL;
	}
	
	//cleanup before writing to pipe
	CloseHandle(PipeInfo->h_pipeRead);
	free(cmd);
	
	if(PipeInfo->Affinity)
		SetProcessAffinityMask(PipeInfo->h_process, PipeInfo->Affinity);
	if(PipeInfo->x264Affinity)
		SetProcessAffinityMask(PipeInfo->pi_info.hProcess, PipeInfo->x264Affinity);
	
	if(VideoInfo->b_mtpipe)
	{
		_beginthread(BufferedPipe, 0, (void*)PipeInfo);
	}
	_beginthread(BufferedReadErrorPipe, 0, (void*)PipeInfo);
}

int WritePipeLoop(int* Func(HANDLE, char *, int), video_info_t * VideoInfo, pipe_info_t * PipeInfo)
{
	unsigned int frame,h_half,w_half;
	AVS_VideoFrame *frm;
	char *planeY, *planeU, *planeV;
	int j;
	//prepare for writing
	w_half = VideoInfo->i_width >> 1;
	h_half = VideoInfo->i_height >> 1;
	int buflength = VideoInfo->i_width * VideoInfo->i_height + h_half * w_half  * 2;
	char title[1024], szbuffer[256];
	if(PipeInfo->BufferSize == 0)
		/* Auto:
			>2M (~1920x1080 = 3.1M) 256x
			>1M (~1280x720 = 1.4M) 512x
			<1M (~848x480 = 610K) 1024x
		*/
		PipeInfo->BufferSize = buflength > 2097152 ? 0x100 : (buflength > 1048576 ? 0x200 : 0x400); 
	
	OBuffer.capacity = PipeInfo->BufferSize;
	OBuffer.dynarray = (node_t*) malloc(sizeof(node_t) * OBuffer.capacity);
	memset(OBuffer.dynarray, 0, sizeof(node_t) * OBuffer.capacity);
	OBuffer.head = OBuffer.tail = 0;
	
	OBuffer.framesize = buflength;
	fprintf( stderr, "avs4x264 [info]: Buffer size set to %dx%d\n", OBuffer.capacity, OBuffer.framesize);
	fflush(stderr);

	//write
	for (frame=PipeInfo->Seek; frame<PipeInfo->Frames; frame++)
	{
		if(VideoInfo->b_mtpipe && (frame & 1))
		{
			GetConsoleTitle(title, 1024);
			if(!strchr(title, '$'))
			{
				sprintf(szbuffer, " $ Buffer: %d/%d (%d MB)", OBuffer.size, OBuffer.capacity, (OBuffer.framesize * OBuffer.size) >> 20);
				strcat(title, szbuffer);
				SetConsoleTitle(title);
			}
		}
		
		frm = avs_h.func.avs_get_frame( avs_h.clip, frame );
		const char *err = avs_h.func.avs_clip_get_error( avs_h.clip );

		if( err )
		{
			fprintf( stderr, "avs [error]: %s occurred while reading frame %d\n", err, frame );
			return ERR_PROCESS_FAIL;
		}
		planeY = (char*)(frm->vfb->data + frm->offset);
		planeU = (char*)(frm->vfb->data + frm->offsetU);
		planeV = (char*)(frm->vfb->data + frm->offsetV);
		
		if(VideoInfo->b_mtpipe)
		{
			char * bufyuv = (char *) malloc(buflength);
			char * bufpos = bufyuv;
			for(j = 0; j < VideoInfo->i_height; j++)
			{
				memcpy(bufpos, planeY, VideoInfo->i_width);
				bufpos += VideoInfo->i_width;
				planeY += frm->pitch;
			}
			for(j = 0; j < h_half; j++)
			{
				memcpy(bufpos, planeU, w_half);
				bufpos += w_half;
				planeU += frm->pitchUV;
			}
			for(j = 0; j < h_half; j++)
			{
				memcpy(bufpos, planeV, w_half);
				bufpos += w_half;
				planeV += frm->pitchUV;
			}
			//assert(bufpos - bufyuv == buflength);
			(*Func)(0, bufyuv, buflength);
		}
		else
		{
			for (j=0; j<VideoInfo->i_height; j++){
				if( !(*Func)(PipeInfo->h_pipeWrite, planeY, VideoInfo->i_width) )
				{
					fprintf( stderr, "avs [error]: Error occurred while writing frame %d\n(Maybe x264_64.exe closed)\n", frame );
					return ERR_PROCESS_FAIL;
				}
				planeY += frm->pitch;
			}
			for (j=0; j<h_half; j++){
				if( !(*Func)(PipeInfo->h_pipeWrite, planeU, w_half) )
				{
					fprintf( stderr, "avs [error]: Error occurred while writing frame %d\n(Maybe x264_64.exe closed)\n", frame );
					return ERR_PROCESS_FAIL;
				}
				planeU += frm->pitchUV;
			}
			for (j=0; j<h_half; j++){
				if( !(*Func)(PipeInfo->h_pipeWrite, planeV, w_half) )
				{
					fprintf( stderr, "avs [error]: Error occurred while writing frame %d\n(Maybe x264_64.exe closed)\n", frame );
					return ERR_PROCESS_FAIL;
				}
				planeV += frm->pitchUV;
			}
		}
		avs_h.func.avs_release_video_frame( frm );
	}
}

